# -*- coding: utf-8 -*-
# @Time : 2021-11-18 16:21
# @Author : lwb
# @Site : 
# @File : transformer.py
import torch.nn as nn
import torch
import torch.nn.functional as F
from utils.Transformer_process import  PositionEncoding,lenght_to_mask
from torch.nn import TransformerEncoderLayer,TransformerEncoder

# 只是用encoder层进行文本分类
class Transformer(nn.Module):
    def __init__(self,vocab_size,embedding_dim,num_class,
                 dim_feedforward=512,num_head=4,num_layers=1,dropout=0.1
                 ,max_len=128,activation:str='relu'):
        super(Transformer, self).__init__()
        self.embedding_dim=embedding_dim
        self.embeddings=nn.Embedding(vocab_size,embedding_dim)                             # 词向量
        self.position_embedding=PositionEncoding(embedding_dim,dropout,max_len)            # 位置编码层
        encoder_layer=TransformerEncoderLayer(embedding_dim,num_head,dim_feedforward,dropout,activation)
        self.transformer=TransformerEncoder(encoder_layer,num_layers)
        self.output=nn.Linear(embedding_dim,num_class)                                         # 输出层

    def forward(self,inputs,lengths):
        inputs=torch.transpose(inputs,0,1)                                                  # 与LSTM处理情况相同，输入数据的第一维是批次，需要转换为所需要的第一维是长度，第二维是批次的形状
        hidden_states=self.embeddings(inputs)
        hidden_states=self.position_embedding(hidden_states)
        attention_mask=lenght_to_mask(lengths)==False                                       # 根据批次中每个序列长度生成Mask矩阵，以便处理长度不一致的序列，忽略掉比较短的序列的无效部分。
        hidden_states=self.transformer(hidden_states,src_key_padding_mask=attention_mask)   # src_key_padding_mask参数正好与length_to_mask函数生成的结果相反（无自注意力部分为True），因此还需要取反，
        hidden_states=hidden_states[0,:,:]                                                  # 取第一个标记的输出结果作为分类层的输入
        output=self.output(hidden_states)
        log_probs=F.log_softmax(output,dim=1)
        return log_probs




